{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JndnmDMp66FL"
      },
      "source": [
        "Copyright 2018 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hMqWDc_m6rUC"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fLgk9Tm3T22v"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import scipy\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "import jax.numpy as jnp\n",
        "import jax\n",
        "import jax.experimental.optimizers\n",
        "import time\n",
        "import functools"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ObB4g9_C3gve"
      },
      "outputs": [],
      "source": [
        "envmap_H = 50\n",
        "envmap_W = 100\n",
        "\n",
        "with gfile.Open(f'{DIRECTORY}/ninomaru_teien_4k.jpg', 'rb') as f:\n",
        "  envmap_gt = np.array(Image.open(f).resize((envmap_W, envmap_H), Image.ANTIALIAS)) / 255.0\n",
        "\n",
        "envmap_gt[0, :, :] = envmap_gt[0, :, :].mean(axis=0, keepdims=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jEiMZYYV3hpd"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=[12, 12])\n",
        "plt.imshow(envmap_gt)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YgLOtJ-l3jvM"
      },
      "outputs": [],
      "source": [
        "phi, theta = jnp.meshgrid(jnp.linspace(0.0, 2*jnp.pi, envmap_W+1)[:-1],\n",
        "                          jnp.linspace(0.0,   jnp.pi, envmap_H+1)[:-1])\n",
        "\n",
        "#mask = np.ones_like(img)\n",
        "\n",
        "x = jnp.sin(theta) * jnp.cos(phi)\n",
        "y = jnp.sin(theta) * jnp.sin(phi)\n",
        "z = jnp.cos(theta)\n",
        "xyz = jnp.stack([x.flatten(), y.flatten(), z.flatten()], axis=-1)\n",
        "\n",
        "def thph2vec(theta, phi):\n",
        "  return jnp.stack([jnp.sin(theta) * jnp.cos(phi),\n",
        "                    jnp.sin(theta) * jnp.sin(phi),\n",
        "                    jnp.cos(theta)], axis=-1)\n",
        "\n",
        "\n",
        "def rotate(xyz, new_z_dir, flipped=True):\n",
        "  new_x_dir = jnp.cross(new_z_dir, jnp.array([0.0, 1.0, 0.0]))\n",
        "  new_x_dir = new_x_dir / jnp.linalg.norm(new_x_dir)\n",
        "  new_y_dir = jnp.cross(new_z_dir, new_x_dir)\n",
        "\n",
        "  R = jnp.stack([new_x_dir, new_y_dir, new_z_dir], axis=1) # columns are x, y, z\n",
        "\n",
        "  if flipped:\n",
        "    new_xyz = xyz @ R\n",
        "  else:\n",
        "    new_xyz = xyz @ R.T\n",
        "\n",
        "  return new_xyz\n",
        "\n",
        "def vec2thph(x, y, z):\n",
        "  theta = jnp.arctan2(jnp.sqrt(x**2+y**2+1e-10), z)\n",
        "  phi = jnp.arctan2(y, x+1e-10)\n",
        "  return theta, phi\n",
        "\n",
        "def get_mask(new_z_theta, new_z_phi, thresh=0.9):\n",
        "\n",
        "  # rotate z -\u003e new_z_dir\n",
        "  new_z_dir = thph2vec(new_z_theta, new_z_phi)\n",
        "  new_xyz = rotate(xyz, new_z_dir)\n",
        "\n",
        "  new_x = new_xyz[:, 0].reshape((envmap_H, envmap_W))\n",
        "  new_y = new_xyz[:, 1].reshape((envmap_H, envmap_W))\n",
        "  new_z = new_xyz[:, 2].reshape((envmap_H, envmap_W))\n",
        "\n",
        "  #new_theta, _ = vec2thph(new_x, new_y, new_z)\n",
        "\n",
        "  #mask = jax.nn.sigmoid(1*(1.0 - new_z))#(new_z \u003c 0.9)\n",
        "  #mask = new_z\n",
        "  mask = jnp.float32(new_z \u003c thresh)\n",
        "  #mask = jax.nn.sigmoid(30.0*(-new_z + thresh))\n",
        "  return mask\n",
        "\n",
        "new_z_theta = jnp.pi/2.0 #jnp.pi/4.0\n",
        "new_z_phi = jnp.pi/2\n",
        "mask = get_mask(new_z_theta, new_z_phi, 0.5)\n",
        "\n",
        "plt.imshow(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u7NePdLY3pjB"
      },
      "outputs": [],
      "source": [
        "def get_rotated_grid(new_z_theta, new_z_phi):\n",
        "  # rotate z -\u003e new_z_dir\n",
        "  new_z_dir = thph2vec(new_z_theta, new_z_phi)\n",
        "  new_xyz = rotate(xyz, new_z_dir)\n",
        "\n",
        "  new_x = new_xyz[:, 0].reshape((envmap_H, envmap_W))\n",
        "  new_y = new_xyz[:, 1].reshape((envmap_H, envmap_W))\n",
        "  new_z = new_xyz[:, 2].reshape((envmap_H, envmap_W))\n",
        "  return new_x, new_y, new_z \n",
        "\n",
        "def get_mask_and_peter(new_z_theta, new_z_phi, thresh=0.9, peter_grid=None, mask_grid=None):\n",
        "  new_x, new_y, new_z = get_rotated_grid(new_z_theta, new_z_phi)\n",
        "\n",
        "  if peter_grid is None:\n",
        "    peter = jnp.zeros((envmap_H, envmap_W, 3), dtype=jnp.float32)\n",
        "    mask = jnp.ones((envmap_H, envmap_W), dtype=jnp.float32)\n",
        "  else:\n",
        "    peter, mask = peter_fn(peter_grid, mask_grid, new_x, new_y, new_z)\n",
        "\n",
        "  return mask, peter, new_x, new_y, new_z\n",
        "\n",
        "def get_illumination(envmap, new_z_theta, new_z_phi, thresh, peter_grid=None, mask_grid=None):\n",
        "  mask, peter, new_x, new_y, new_z = get_mask_and_peter(new_z_theta, new_z_phi, thresh, peter_grid, mask_grid)\n",
        "  return mask[:, :, None] * envmap + (1.0 - mask[:, :, None]) * peter, mask, peter\n",
        "\n",
        "\n",
        "def peter_fn(peter_grid, mask_grid, x, y, z):\n",
        "  th, ph = vec2thph(x, y, z)\n",
        "  inds = jnp.stack([th*envmap_H/jnp.pi, ph*envmap_W/2.0/jnp.pi], axis=0)\n",
        "  peter = sphere_interpolate(peter_grid, inds)\n",
        "  mask = sphere_interpolate(mask_grid[..., None], inds)[..., 0]\n",
        "  return peter, mask\n",
        "\n",
        "def sphere_interpolate(grids, inds):\n",
        "  \"\"\"\n",
        "  grid is [H, W, d], parameterized by elevation (0th dim) and azimuth (1st dim)\n",
        "  inds is [2, ...], with the 0th dim being elevation and 1st azimuth\n",
        "  \"\"\"\n",
        "\n",
        "  #inds[1, ...] = jnp.mod(inds[1, ...], 2*jnp.pi)\n",
        "  results = []\n",
        "  for grid in [grids[:, :, d] for d in range(grids.shape[-1])]:\n",
        "    res = jax.scipy.ndimage.map_coordinates(grid, inds, order=1, mode='nearest')\n",
        "    #res = jax.scipy.ndimage.map_coordinates(grid, inds, order=1, mode='wrap')\n",
        "    results.append(res)\n",
        "\n",
        "  return jnp.stack(results, axis=-1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p9Vf02so3rar"
      },
      "outputs": [],
      "source": [
        "def cosine_scheduler(m, lr_init, lr_final):\n",
        "  def lr(i):\n",
        "    return 0.5 * (1.0 - jnp.cos(jnp.pi*i/m)) * (lr_final - lr_init) + lr_init\n",
        "  return lr\n",
        "\n",
        "def parse_bin(s):\n",
        "  return int(s[1:], 2) / 2.**(len(s) - 1)\n",
        "\n",
        "\n",
        "def phi2(i):\n",
        "  return parse_bin('.' + f'{i:b}'[::-1])\n",
        "\n",
        "def nice_uniform(N):\n",
        "  u = []\n",
        "  v = []\n",
        "  for i in range(N):\n",
        "      u.append(i / float(N))\n",
        "      v.append(phi2(i))\n",
        "      #pts.append((i/float(N), phi2(i)))\n",
        "\n",
        "  return u, v\n",
        "\n",
        "def nice_uniform_hemispherical(N):\n",
        "  \"\"\"implementation of http://holger.dammertz.org/stuff/notes_HammersleyOnHemisphere.html\"\"\"\n",
        "  u, v = nice_uniform(N)\n",
        "\n",
        "  theta = np.arccos(1.0 - np.array(u))\n",
        "  phi   = 2.0 * np.pi * np.array(v)\n",
        "\n",
        "  return np.stack([theta, phi], axis=-1)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JklMXS5o3tc_"
      },
      "outputs": [],
      "source": [
        "thresh = 0.85\n",
        "\n",
        "peter_grid_gt = np.stack([x, y, z], axis=-1) * 0.5 + 0.5\n",
        "peter_grid_gt *= 0.0\n",
        "\n",
        "mask_grid_gt = np.float32(np.cos(theta) \u003c thresh)\n",
        "\n",
        "\n",
        "num_images = 1000\n",
        "envmaps_gt = []\n",
        "masks_gt = []\n",
        "peters_gt = []\n",
        "mask_thetaphis = []\n",
        "for i, (mask_theta, mask_phi) in enumerate(nice_uniform_hemispherical(num_images)):\n",
        "  envmap, mask, peter = get_illumination(envmap_gt, mask_theta, mask_phi, thresh, peter_grid_gt, mask_grid_gt)\n",
        "  envmaps_gt.append(envmap)\n",
        "  masks_gt.append(mask)\n",
        "  peters_gt.append(peter)\n",
        "  mask_thetaphis.append((mask_theta, mask_phi))\n",
        "\n",
        "\n",
        "#print(len(mask_phis))\n",
        "plt.figure()\n",
        "plt.plot(*zip(*mask_thetaphis), '.')  \n",
        "\n",
        "masks_gt = jnp.array(masks_gt)\n",
        "peters_gt = jnp.array(peters_gt)\n",
        "envmaps_gt = jnp.array(envmaps_gt)\n",
        "orientations_gt = jnp.array(mask_thetaphis)\n",
        "#valid = 1.0 - jnp.float32(invalid).flatten()[None, None, :]\n",
        "M = masks_gt.shape[0]\n",
        "\n",
        "\n",
        "\n",
        "plt.figure()\n",
        "plt.imshow(masks_gt.min(axis=0))\n",
        "plt.figure()\n",
        "plt.imshow((1.0 - masks_gt).sum(0))\n",
        "plt.colorbar()\n",
        "\n",
        "\n",
        "H = W = 96\n",
        "\n",
        "xx, yy = np.meshgrid(np.linspace(-1.01, 1.01, W),\n",
        "                     np.linspace(-1.01, 1.01, H))\n",
        "\n",
        "invalid = xx ** 2 + yy ** 2 \u003e 1.0\n",
        "\n",
        "zz = np.sqrt(1.0 - xx ** 2 - yy ** 2)\n",
        "\n",
        "xx[invalid] = 0.0\n",
        "yy[invalid] = 0.0\n",
        "zz[invalid] = 1.0\n",
        "\n",
        "normals_gt = jnp.stack([xx.flatten(), yy.flatten(), zz.flatten()], axis=-1)\n",
        "\n",
        "plt.figure()\n",
        "plt.imshow(normals_gt.reshape(H, W, 3)*0.5+0.5)\n",
        "\n",
        "num_eqs  = (1 - invalid).sum() * masks_gt.shape[0]\n",
        "num_vars = envmap_H * envmap_W * 3 + masks_gt.shape[0] * envmap_H * envmap_W\n",
        "\n",
        "#assert (1 - masks_gt).sum(0).min() \u003e= 1, 'Some pixels are not covered by any mask'\n",
        "assert num_eqs \u003e num_vars, f\"Need enough pixels ({num_eqs} eqs, {num_vars} vars)\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-YxTulFl3zeO"
      },
      "outputs": [],
      "source": [
        "plt.imshow(get_mask_and_peter(1, 4.0, peter_grid=peter_grid_gt, mask_grid=mask_grid_gt)[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aDjEk09i4my-"
      },
      "outputs": [],
      "source": [
        "@jax.custom_jvp\n",
        "def render_pixels(normals, envmap, peter, mask, use_jacobian=True):\n",
        "  \"\"\"\n",
        "  normals:      [3]\n",
        "  envmap:       [H, W, 3]\n",
        "  peter:        [H, W, 3]  (grid, interpolated using `peter_fn`)\n",
        "  mask:         [H, W]\n",
        "\n",
        "  returns [3]\n",
        "  \"\"\"\n",
        "  # xyz is [HW, 3]\n",
        "  lobe = jnp.maximum(0.0, (xyz * normals[None, :]).sum(-1).reshape(envmap_H, envmap_W))  # [H, W, 1]\n",
        "  #print(normals.shape, envmap.shape, peter.shape, mask.shape) (3,) (100, 200, 3) (200, 3) (200,)\n",
        "\n",
        "  masked_envmap = envmap * mask[:, :, None] + (1.0 - mask[:, :, None]) * peter\n",
        "  if use_jacobian:\n",
        "    return (masked_envmap * lobe[:, :, None] * jnp.sin(theta)[:, :, None]).sum(0).sum(0) * (phi[0, 1] - phi[0, 0]) * (theta[1, 0] - theta[0, 0]) / jnp.pi\n",
        "  else:\n",
        "    return (masked_envmap * lobe[:, :, None]).sum(0).sum(0) * (phi[0, 1] - phi[0, 0]) * (theta[1, 0] - theta[0, 0]) / jnp.pi\n",
        "\n",
        "\n",
        "@render_pixels.defjvp\n",
        "def render_pixels_jvp(primals, tangents):\n",
        "  a, b, x, filt = primals\n",
        "  a_dot, b_dot, x_dot, filt_dot = tangents\n",
        "\n",
        "  primal_out = params2signal(a, b, x, filt)\n",
        "  xa = periodic_interp(a, t, x)\n",
        "  xb = periodic_interp(b, t, x)\n",
        "\n",
        "  ftma = periodic_interp(t - a, t, filt)\n",
        "  ftmb = periodic_interp(t - b, t, filt)\n",
        "\n",
        "  # TODO: figure out how to reuse masks from the main function instead of\n",
        "  #       recomputing them here.\n",
        "  mask = (b \u003e a) * (t \u003e a) * (t \u003c b) \\\n",
        "        + (1.0 - (b \u003e a)) * (1.0 - (t \u003e b) * (t \u003c a))\n",
        "\n",
        "  tangent_out = -xa * ftma * a_dot + xb * ftmb * b_dot + \\\n",
        "    jnp.real(jnp.fft.ifft(jnp.fft.fft(filt) * jnp.fft.fft(mask * x_dot, axis=-1), axis=-1)) + \\\n",
        "    jnp.real(jnp.fft.ifft(jnp.fft.fft(filt_dot) * jnp.fft.fft(mask * x, axis=-1), axis=-1))\n",
        "  return primal_out, tangent_out\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def make_data(envmap, peters, masks, normals):\n",
        "  \"\"\"\n",
        "  normals:       [   n,       3]\n",
        "  envmap:        [      H, W, 3]\n",
        "  peters:        [N,    H, W, 3]\n",
        "  masks:         [N,    H, W   ]\n",
        "  \"\"\"\n",
        "  return jax.vmap(jax.vmap(render_pixels, in_axes=(0, None, None, None)), in_axes=(None, None, 0, 0))(normals, envmap, peters, masks)\n",
        "\n",
        "\n",
        "def get_diff(params, num_pts=20.0):\n",
        "  \"\"\"\n",
        "  Compute the gradient of the image w.r.t. the vertex parameters\n",
        "  \"\"\"\n",
        "  t = (jnp.arange(num_pts) + 0.5) / num_pts\n",
        "  r = params[:, :, None].transpose(1, 0, 2) * t[None, None, :] + jnp.roll(params, 1, axis=0)[:, :, None].transpose(1, 0, 2) * (1.0 - t[None, None, :])\n",
        "\n",
        "  fq = (0.5 * jnp.stack([x, y], axis=0)[:, None, None, :, :] + 0.5) * (H - 1) - r[:, :, :, None, None]\n",
        "\n",
        "  f_interp = interp2d(f[:, :, None], fq)\n",
        "  L_interp = interp2d(L[:, :, None], r) \n",
        "\n",
        "  line_diff1 = (f_interp[:, :, :, :, 0] * L_interp[:, :, None, None, 0] * (1.0 - t[None, :, None, None])).sum(1) * (t[1] - t[0])  # the ith row is the integral from v_{i} to v_{i-1}\n",
        "  line_diff2 = (f_interp[:, :, :, :, 0] * L_interp[:, :, None, None, 0] * t[None, :, None, None]).sum(1) * (t[1] - t[0])          # the ith row is the integral from v_{i-1} to v_{i}\n",
        "  dx = jnp.fliplr(params - jnp.roll(params, 1, axis=0)) * jnp.array([1.0, -1.0])[None, :]  # this stores y_{i} - y_{i-1} and x_{i} - x_{i-1})\n",
        "\n",
        "  # Add integral from v_{i+1} to v_{i} multiplied by y_{i} - y_{i+1} with integral from v_{i-1} to v_{i} multiplied by y_{i-1} - y_{i}\n",
        "  diff = -jnp.roll(line_diff1[:, None, :, :] * dx[:, :, None, None], -1, axis=0) - line_diff2[:, None, :, :] * dx[:, :, None, None]\n",
        "  return diff\n",
        "\n",
        "\n",
        "def get_mask(vertices):\n",
        "  img = np.ones((H, W), dtype=np.uint8) * 255\n",
        "  mask = cv2.fillPoly(img, vertices, 0, cv2.LINE_AA)\n",
        "  return np.float32(mask) / 255.0\n",
        "\n",
        "\"\"\"\n",
        "@jax.custom_jvp\n",
        "def params_to_data(envmap, peters, vertices, normals):\n",
        "  \n",
        "  #a, shape ()\n",
        "  #b, shape ()\n",
        "  #x, shape (T,)\n",
        "  #filt, shape (T,)\n",
        "  #\n",
        "  #t is not an input but it also has shape (T,) \n",
        "  \n",
        "  mask = get_mask(vertices)\n",
        "\n",
        "  #shifted_peters = jax.vmap(periodic_interp, in_axes=(0, None, None))(t[None, :]-positions[:, None], t, peter)\n",
        "  peter = jnp.zeros_like(mask)\n",
        "  print(\"NO PETER!\")\n",
        "  \n",
        "  y = mask * x + (1.0 - mask) * peter\n",
        "  z = jnp.real(jnp.fft.ifft(jnp.fft.fft(filt) * jnp.fft.fft(y, axis=-1), axis=-1))\n",
        "\n",
        "  return z\n",
        "\n",
        "@params2signal.defjvp\n",
        "def params2signal_jvp(primals, tangents):\n",
        "  a, b, x, filt = primals\n",
        "  a_dot, b_dot, x_dot, filt_dot = tangents\n",
        "\n",
        "  primal_out = params2signal(a, b, x, filt)\n",
        "  xa = periodic_interp(a, t, x)\n",
        "  xb = periodic_interp(b, t, x)\n",
        "\n",
        "  ftma = periodic_interp(t - a, t, filt)\n",
        "  ftmb = periodic_interp(t - b, t, filt)\n",
        "\n",
        "  # TODO: figure out how to reuse masks from the main function instead of\n",
        "  #       recomputing them here.\n",
        "  mask = (b \u003e a) * (t \u003e a) * (t \u003c b) \\\n",
        "        + (1.0 - (b \u003e a)) * (1.0 - (t \u003e b) * (t \u003c a))\n",
        "\n",
        "  tangent_out = -xa * ftma * a_dot + xb * ftmb * b_dot + \\\n",
        "    jnp.real(jnp.fft.ifft(jnp.fft.fft(filt) * jnp.fft.fft(mask * x_dot, axis=-1), axis=-1)) + \\\n",
        "    jnp.real(jnp.fft.ifft(jnp.fft.fft(filt_dot) * jnp.fft.fft(mask * x, axis=-1), axis=-1))\n",
        "  return primal_out, tangent_out\n",
        "\"\"\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yWpb3CB34oBm"
      },
      "outputs": [],
      "source": [
        "res_gt = []\n",
        "num_chunks = 64\n",
        "assert normals_gt.shape[0] % num_chunks == 0\n",
        "for chunk in range(num_chunks):\n",
        "  ri = chunk * normals_gt.shape[0] // num_chunks\n",
        "  rf = ri    + normals_gt.shape[0] // num_chunks\n",
        "  res_gt.append(make_data(envmap_gt, peters_gt, masks_gt, normals_gt[ri:rf, :]))\n",
        "res_gt = jnp.concatenate(res_gt, axis=1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XXgB7maH4pTM"
      },
      "outputs": [],
      "source": [
        "def entropy(p):\n",
        "  return -p * jnp.log(jnp.maximum(1e-10, p)) - (1 - p) * jnp.log(jnp.maximum(1e-10, 1.0 - p))\n",
        "\n",
        "\n",
        "def norm(x, p=1, w=None):\n",
        "  if w is not None:\n",
        "    x = x * w\n",
        "  if p == 0.0:\n",
        "    return jnp.float32(jnp.sum(x != 0))\n",
        "  return jnp.power((jnp.abs(x + 1e-3) ** p).sum(), 1.0/p)\n",
        "\n",
        "@jax.jit\n",
        "def get_loss(params_envmap, params_peters, params_masks, normals, gt, spatial_inds, img_inds, i):\n",
        "  \"\"\"\n",
        "  envmap is  [   H, W, 3]\n",
        "  peters is  [N, H, W, 3]\n",
        "  masks is   [N, H, W]\n",
        "  normals is [N,       3]\n",
        "  gt is      [N, H*W,  3]\n",
        "  \"\"\"\n",
        "  envmap, peters, masks = params2components(params_envmap, params_peters, params_masks, i)\n",
        "\n",
        "  res = make_data(envmap, peters[img_inds], masks[img_inds], normals[spatial_inds])\n",
        "\n",
        "  data_loss = jnp.sum(((gt[img_inds, :, :][:, spatial_inds, :] - res)**2)) / batch_size #/ gt.shape[0]\n",
        "  loss = data_loss#.clone()\n",
        "\n",
        "  # TV norm is the integral of gradient(mask) * sin(theta) dtheta dphi\n",
        "  #p = 0.1\n",
        "  w = masks.shape[0] * masks.shape[1] * masks.shape[2]\n",
        "  #mask_tv = norm(masks[:, 1:, :] - masks[:, :-1, :], p, w=jnp.sin(theta[None, :-1, :])) / w + norm(masks[:, :, 1:] - masks[:, :, :-1], p) / w\n",
        "  \n",
        "  \n",
        "  #norm = lambda x: lossfun(x, -cosine_scheduler(num_iters, 2.0, -10.0)(i), 0.1)\n",
        "  #norm = lambda x: lossfun(x, -1.0, 0.1)\n",
        "  norm = lambda x: x ** 2\n",
        "  sintheta = jnp.sin(theta)\n",
        "  mask_tv = norm((masks[:, 1:, :] - masks[:, :-1, :])*sintheta[None, :-1, :]).sum() + norm((masks[:, :, 1:] - masks[:, :, :-1])*sintheta[None, :, :-1]).sum()\n",
        "  # Add wrap gradient in phi\n",
        "  mask_tv += norm((masks[:, :, :1] - masks[:, :, -1:])*sintheta[None, :, -1:]).sum()\n",
        "  mask_tv *= (phi[0, 1] - phi[0, 0]) * (theta[1, 0] - theta[0, 0])\n",
        "  loss += 1e-4 * mask_tv #/ w\n",
        "  #loss += 0.01 * (1.0 - masks).sum() / w\n",
        "  loss += 1e-2 * entropy(masks).sum() / w\n",
        "\n",
        "  #loss += 1e-5 * ((peters - peters.mean(0, keepdims=True)) ** 2).sum() / peters.shape[0]\n",
        "\n",
        "  return loss, (data_loss, None, res, ((envmap-envmap_gt)**2).mean())\n",
        "  \n",
        "\n",
        "\n",
        "\n",
        "\n",
        "params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_W, 3)) - 0.5) * 0.1\n",
        "#params_peters = (jax.random.uniform(jax.random.PRNGKey(7065), shape=(res_gt.shape[0], envmap_H, envmap_W, 3)) - 0.5) * 0.1 #- 3.0\n",
        "params_peters = (jax.random.uniform(jax.random.PRNGKey(7065), shape=(res_gt.shape[0], 1, 1, 3)) - 0.5) * 0.1\n",
        "#params_peters = (jax.random.uniform(jax.random.PRNGKey(7065), shape=(1, 1, 1, 3)) - 0.5) * 0.1\n",
        "if True:\n",
        "  #params_masks =  (jax.random.uniform(jax.random.PRNGKey(1122), shape=(res_gt.shape[0], envmap_H, envmap_W)) - 0.5) * 10.0 + 3.0\n",
        "  #params_masks =  jax.random.uniform(jax.random.PRNGKey(1122), shape=(res_gt.shape[0], envmap_H, envmap_W)) * 1.0 #+ 1.0\n",
        "  params_masks =  jax.random.uniform(jax.random.PRNGKey(1122), shape=(res_gt.shape[0], envmap_H, envmap_W)) * 1.0 + 1.0\n",
        "else:\n",
        "  params_masks = jax.random.uniform(jax.random.PRNGKey(1122), shape=(res_gt.shape[0], 256, 2))\n",
        "\n",
        "\n",
        "#params_envmap = envmap_gt\n",
        "#params_peters = peters_gt\n",
        "#params_masks = masks_gt\n",
        "\n",
        "jitvalgrad = jax.jit(jax.value_and_grad(get_loss, argnums=(0, 1, 2), has_aux=True))\n",
        "\n",
        "num_iters = 50000\n",
        "spatial_batch_size = 4\n",
        "img_batch_size = 128\n",
        "batch_size = spatial_batch_size * img_batch_size\n",
        "print(f\"batch_size is {batch_size}\")\n",
        "\n",
        "\n",
        "init_lr_envmap = 0.003\n",
        "init_envmap, update_envmap, get_params_envmap = jax.experimental.optimizers.adam(cosine_scheduler(num_iters, init_lr_envmap, init_lr_envmap*0.3))\n",
        "state_envmap = init_envmap(params_envmap)\n",
        "\n",
        "init_lr_peters = 0.01\n",
        "init_peters, update_peters, get_params_peters = jax.experimental.optimizers.adam(cosine_scheduler(num_iters, init_lr_peters, init_lr_peters*0.3))\n",
        "state_peters = init_peters(params_peters)\n",
        "\n",
        "init_lr_masks = 0.07\n",
        "init_masks, update_masks, get_params_masks = jax.experimental.optimizers.adam(cosine_scheduler(num_iters, init_lr_masks, init_lr_masks*0.3))\n",
        "state_masks = init_masks(params_masks)\n",
        "\n",
        "def params2components(params_envmap, params_peters, params_masks, iteration):\n",
        "  envmap = jax.nn.sigmoid(params_envmap)# / jnp.sin(theta[:, :, None] + 1e-10))\n",
        "  peters = jax.nn.sigmoid(params_peters) #* 0.0\n",
        "  if True:\n",
        "    #temperature = cosine_scheduler(num_iters, 1e-1, 1e-4)(iteration)\n",
        "    temperature = 1.0\n",
        "    masks = jax.nn.sigmoid(params_masks/temperature) #/ jnp.sin(theta[None, :, :] + 1e-10) #* 0.0 + 1.0\n",
        "  else:\n",
        "    width = 0.3\n",
        "    distsq = (params_masks[:, :, 0][:, :, None, None] * jnp.pi - theta[None, None, :, :]) ** 2 + (params_masks[:, :, 1][:, :, None, None] * 2 * jnp.pi - phi[None, None, :, :]) ** 2\n",
        "    masks = jnp.exp(-distsq/2/width).sum(1)\n",
        "    masks = jnp.clip(masks, 0.0, 1.0)\n",
        "  \n",
        "  if True:\n",
        "    #envmap = envmap_gt\n",
        "    peters = peters_gt\n",
        "    masks = masks_gt\n",
        "    for _ in range(20):\n",
        "      print(\"SETTING STUFF TO GT\")\n",
        "\n",
        "\n",
        "  return envmap, peters, masks\n",
        "\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def update_params(i, state_envmap, state_peters, state_masks, normals, gt, spatial_inds, img_inds):\n",
        "  params_envmap = get_params_envmap(state_envmap)\n",
        "  params_peters = get_params_peters(state_peters)\n",
        "  params_masks  = get_params_masks(state_masks)\n",
        "\n",
        "  (loss, (data_loss, _, res, envmap_error)), g = jitvalgrad(params_envmap, params_peters, params_masks, normals, gt, spatial_inds, img_inds, i)\n",
        "\n",
        "  return update_envmap(i, g[0], state_envmap), update_peters(i, g[1], state_peters), update_masks(i, g[2], state_masks), loss, data_loss, res, envmap_error, g\n",
        "\n",
        "from numpy.random import default_rng\n",
        "\n",
        "rng = default_rng()\n",
        "\n",
        "valid = 1.0 - jnp.float32(invalid).flatten()[None, None, :]\n",
        "nz = np.nonzero(valid[0, 0, :])[0]  # valid element indices\n",
        "\n",
        "\n",
        "\n",
        "t = time.time()\n",
        "losses = []\n",
        "envmap_errors = []\n",
        "for i in range(num_iters):\n",
        "  if i % 100 == 0:\n",
        "    print(i)\n",
        "  spatial_inds = nz[rng.choice(nz.shape[0], size=spatial_batch_size, replace=False)]\n",
        "  img_inds = rng.choice(res_gt.shape[0], size=img_batch_size, replace=False)\n",
        "  state_envmap, state_peters, state_masks, loss, data_loss, z, envmap_error, g = update_params(\n",
        "      i, state_envmap, state_peters, state_masks, normals_gt, res_gt, spatial_inds, img_inds)\n",
        "  losses.append(data_loss)\n",
        "  envmap_errors.append(envmap_error)\n",
        "\n",
        "  if (i % 5000 == 0 and i \u003e 0) or i == num_iters - 1 or i == 2000:\n",
        "\n",
        "    envmap, peters, masks = params2components(get_params_envmap(state_envmap), get_params_peters(state_peters), get_params_masks(state_masks), i)\n",
        "\n",
        "    print(f\"Iteration {i}\")\n",
        "    plt.figure()\n",
        "    plt.semilogy(np.array(losses))\n",
        "    plt.semilogy(np.array(envmap_errors))\n",
        "    plt.figure(figsize=[18, 12])\n",
        "    plt.subplot(321)\n",
        "    plt.imshow(envmap) \n",
        "    plt.axis('off')\n",
        "    plt.subplot(322)\n",
        "    plt.imshow(envmap_gt)\n",
        "    plt.axis('off')\n",
        "    for ind, mask_ind in zip([3, 4, 5], [0, 30, 60]):\n",
        "      plt.subplot(3, 2, ind)\n",
        "      mask_vs_gt = jnp.stack([masks[mask_ind], masks_gt[mask_ind], jnp.ones_like(masks[0])], axis=-1)\n",
        "      plt.imshow(mask_vs_gt)\n",
        "      plt.axis('off')\n",
        "    plt.subplot(326)\n",
        "    #plt.imshow(jnp.where(theta[:, :, None] \u003e= np.arccos(thresh), np.ones_like(peter_grid_gt)*0.5, peter_grid_gt))\n",
        "    plt.imshow(peters[:, 0, :, :].repeat(peters.shape[0], 1))\n",
        "    plt.axis('off')\n",
        "    plt.show()\n",
        "    #print(peters[0, 0, 0, :])\n",
        "\n",
        "print(f\"Optimization took {time.time()-t:.3f} seconds.\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tCi_eb_x45SI"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/dm_python:dm_notebook3",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1BJqgNvhS-KwEwCTgl_QEUrd3Aoa07K2U",
          "timestamp": 1667588239785
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
